{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VZ-KA8k5kybx"
      },
      "source": [
        "# Introduction to tensor slicing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/tensor_slicing\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/tensor_slicing.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/tensor_slicing.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/tensor_slicing.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AixIdVeRk3CO"
      },
      "source": [
        "When working on ML applications such as object detection and NLP, it is sometimes necessary to work with sub-sections (slices) of tensors. For example, if your model architecture includes routing, where one layer might control which training example gets routed to the next layer. In this case, you could use tensor slicing ops to split the tensors up and put them back together in the right order.\n",
        "\n",
        "In NLP applications, you can use tensor slicing to perform word masking while training. For example, you can generate training data from a list of sentences by choosing a word index to mask in each sentence, taking the word out as a label, and then replacing the chosen word with a mask token.  \n",
        "\n",
        "In this guide, you will learn how to use the TensorFlow APIs to:\n",
        "\n",
        "* Extract slices from a tensor\n",
        "* Insert data at specific indices in a tensor\n",
        "\n",
        "This guide assumes familiarity with tensor indexing. Read the indexing sections of the [Tensor](https://www.tensorflow.org/guide/tensor#indexing) and [TensorFlow NumPy](https://www.tensorflow.org/guide/tf_numpy#indexing) guides before getting started with this guide."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FcWhWYn7eXkF"
      },
      "source": [
        "## Setup\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m6uvewqi0jso"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K-muS4ej5zoN"
      },
      "source": [
        "## Extract tensor slices\n",
        "\n",
        "Perform NumPy-like tensor slicing using `tf.slice`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wZep0cjs0Oai"
      },
      "outputs": [],
      "source": [
        "t1 = tf.constant([0, 1, 2, 3, 4, 5, 6, 7])\n",
        "\n",
        "print(tf.slice(t1,\n",
        "               begin=[1],\n",
        "               size=[3]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vh3xI3j0DRJ2"
      },
      "source": [
        "Alternatively, you can use a more Pythonic syntax. Note that tensor slices are evenly spaced over a start-stop range."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P1MtEyKuWuDD"
      },
      "outputs": [],
      "source": [
        "print(t1[1:4])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cjq1o8D2wKKs"
      },
      "source": [
        "<img src=\"images/tf_slicing/slice_1d_1.png\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UunuLTIuwDA-"
      },
      "outputs": [],
      "source": [
        "print(t1[-3:])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EHvRB-XTwRTd"
      },
      "source": [
        "<img src=\"images/tf_slicing/slice_1d_2.png\">"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SW1zFFTnUpCQ"
      },
      "source": [
        "For 2-dimensional tensors,you can use something like:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kThZhmpAVAQw"
      },
      "outputs": [],
      "source": [
        "t2 = tf.constant([[0, 1, 2, 3, 4],\n",
        "                  [5, 6, 7, 8, 9],\n",
        "                  [10, 11, 12, 13, 14],\n",
        "                  [15, 16, 17, 18, 19]])\n",
        "\n",
        "print(t2[:-1, 1:3])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xA5Xt4OdVUui"
      },
      "source": [
        "<img src=\"images/tf_slicing/slice_2d_1.png\">"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iJPggqsH15fI"
      },
      "source": [
        "You can use `tf.slice` on higher dimensional tensors as well."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Re5eX1OXnKOZ"
      },
      "outputs": [],
      "source": [
        "t3 = tf.constant([[[1, 3, 5, 7],\n",
        "                   [9, 11, 13, 15]],\n",
        "                  [[17, 19, 21, 23],\n",
        "                   [25, 27, 29, 31]]\n",
        "                  ])\n",
        "\n",
        "print(tf.slice(t3,\n",
        "               begin=[1, 1, 0],\n",
        "               size=[1, 1, 2]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x-O5FNV9qOJK"
      },
      "source": [
        "You can also use `tf.strided_slice` to extract slices of tensors by 'striding' over the tensor dimensions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b9FhvrOnJsJb"
      },
      "source": [
        "Use `tf.gather` to extract specific indices from a single axis of a tensor."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TwviZrrIj2h7"
      },
      "outputs": [],
      "source": [
        "print(tf.gather(t1,\n",
        "                indices=[0, 3, 6]))\n",
        "\n",
        "# This is similar to doing\n",
        "\n",
        "t1[::3]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oKyjGi2zyzEC"
      },
      "source": [
        "<img src=\"images/tf_slicing/slice_1d_3.png\">"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "obrjeKy1WfTN"
      },
      "source": [
        "`tf.gather` does not require indices to be evenly spaced."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LjJcwcZ0druw"
      },
      "outputs": [],
      "source": [
        "alphabet = tf.constant(list('abcdefghijklmnopqrstuvwxyz'))\n",
        "\n",
        "print(tf.gather(alphabet,\n",
        "                indices=[2, 0, 19, 18]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mSHmUXIyeaJG"
      },
      "source": [
        "<img src=\"images/tf_slicing/gather_1.png\">"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XsxMx49SOaVu"
      },
      "source": [
        "To extract slices from multiple axes of a tensor, use `tf.gather_nd`. This is useful when you want to gather the elements of a matrix as opposed to just its rows or columns."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mT52NFWVdiTe"
      },
      "outputs": [],
      "source": [
        "t4 = tf.constant([[0, 5],\n",
        "                  [1, 6],\n",
        "                  [2, 7],\n",
        "                  [3, 8],\n",
        "                  [4, 9]])\n",
        "\n",
        "print(tf.gather_nd(t4,\n",
        "                   indices=[[2], [3], [0]]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "87NN7YQhh2-a"
      },
      "source": [
        "<img src=\"images/tf_slicing/gather_2.png\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_z6F2WcPJ9Rh"
      },
      "outputs": [],
      "source": [
        "t5 = np.reshape(np.arange(18), [2, 3, 3])\n",
        "\n",
        "print(tf.gather_nd(t5,\n",
        "                   indices=[[0, 0, 0], [1, 2, 1]]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gyIjhm7cV2N0"
      },
      "outputs": [],
      "source": [
        "# Return a list of two matrices\n",
        "\n",
        "print(tf.gather_nd(t5,\n",
        "                   indices=[[[0, 0], [0, 2]], [[1, 0], [1, 2]]]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "368D4ciDWB3r"
      },
      "outputs": [],
      "source": [
        "# Return one matrix\n",
        "\n",
        "print(tf.gather_nd(t5,\n",
        "                   indices=[[0, 0], [0, 2], [1, 0], [1, 2]]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "od51VzS2SSPS"
      },
      "source": [
        "## Insert data into tensors\n",
        "\n",
        "Use `tf.scatter_nd` to insert data at specific slices/indices of a tensor. Note that the tensor into which you insert values is zero-initialized."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jlALYLWm1KhN"
      },
      "outputs": [],
      "source": [
        "t6 = tf.constant([10])\n",
        "indices = tf.constant([[1], [3], [5], [7], [9]])\n",
        "data = tf.constant([2, 4, 6, 8, 10])\n",
        "\n",
        "print(tf.scatter_nd(indices=indices,\n",
        "                    updates=data,\n",
        "                    shape=t6))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CD5vd-kxksW7"
      },
      "source": [
        "Methods like `tf.scatter_nd` which require zero-initialized tensors are similar to sparse tensor initializers. You can use `tf.gather_nd` and `tf.scatter_nd` to mimic the behavior of sparse tensor ops.\n",
        "\n",
        "Consider an example where you construct a sparse tensor using these two methods in conjunction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xyK69QgRmrlW"
      },
      "outputs": [],
      "source": [
        "# Gather values from one tensor by specifying indices\n",
        "\n",
        "new_indices = tf.constant([[0, 2], [2, 1], [3, 3]])\n",
        "t7 = tf.gather_nd(t2, indices=new_indices)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_7V_Qfa4qkdn"
      },
      "source": [
        "<img src=\"images/tf_slicing/gather_nd_sparse.png\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QWT1E1Weqjx2"
      },
      "outputs": [],
      "source": [
        "# Add these values into a new tensor\n",
        "\n",
        "t8 = tf.scatter_nd(indices=new_indices, updates=t7, shape=tf.constant([4, 5]))\n",
        "\n",
        "print(t8)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NUyYjnvCn_vu"
      },
      "source": [
        "This is similar to:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LeqFwUgroE4j"
      },
      "outputs": [],
      "source": [
        "t9 = tf.SparseTensor(indices=[[0, 2], [2, 1], [3, 3]],\n",
        "                     values=[2, 11, 18],\n",
        "                     dense_shape=[4, 5])\n",
        "\n",
        "print(t9)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5MaF6RlJot33"
      },
      "outputs": [],
      "source": [
        "# Convert the sparse tensor into a dense tensor\n",
        "\n",
        "t10 = tf.sparse.to_dense(t9)\n",
        "\n",
        "print(t10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4sf3F3Xk56Bt"
      },
      "source": [
        "To insert data into a tensor with pre-existing values, use `tf.tensor_scatter_nd_add`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mte2ifOb6sQO"
      },
      "outputs": [],
      "source": [
        "t11 = tf.constant([[2, 7, 0],\n",
        "                   [9, 0, 1],\n",
        "                   [0, 3, 8]])\n",
        "\n",
        "# Convert the tensor into a magic square by inserting numbers at appropriate indices\n",
        "\n",
        "t12 = tf.tensor_scatter_nd_add(t11,\n",
        "                               indices=[[0, 2], [1, 1], [2, 0]],\n",
        "                               updates=[6, 5, 4])\n",
        "\n",
        "print(t12)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2dQYyROU09G6"
      },
      "source": [
        "Similarly, use `tf.tensor_scatter_nd_sub` to subtract values from a tensor with pre-existing values."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ac6_i6uK1EI6"
      },
      "outputs": [],
      "source": [
        "# Convert the tensor into an identity matrix\n",
        "\n",
        "t13 = tf.tensor_scatter_nd_sub(t11,\n",
        "                               indices=[[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [2, 1], [2, 2]],\n",
        "                               updates=[1, 7, 9, -1, 1, 3, 7])\n",
        "\n",
        "print(t13)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B_2DuzRRwVc8"
      },
      "source": [
        "Use `tf.tensor_scatter_nd_min` to copy element-wise minimum values from one tensor to another."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T_4FrHrHlkHK"
      },
      "outputs": [],
      "source": [
        "t14 = tf.constant([[-2, -7, 0],\n",
        "                   [-9, 0, 1],\n",
        "                   [0, -3, -8]])\n",
        "\n",
        "t15 = tf.tensor_scatter_nd_min(t14,\n",
        "                               indices=[[0, 2], [1, 1], [2, 0]],\n",
        "                               updates=[-6, -5, -4])\n",
        "\n",
        "print(t15)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PkaiKyrF0WtX"
      },
      "source": [
        "Similarly, use `tf.tensor_scatter_nd_max` to copy element-wise maximum values from one tensor to another."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "izJu0nXi0GDq"
      },
      "outputs": [],
      "source": [
        "t16 = tf.tensor_scatter_nd_max(t14,\n",
        "                               indices=[[0, 2], [1, 1], [2, 0]],\n",
        "                               updates=[6, 5, 4])\n",
        "\n",
        "print(t16)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QAffUOa-85lF"
      },
      "source": [
        "## Further reading and resources\n",
        "\n",
        "In this guide, you learned how to use the tensor slicing ops available with TensorFlow to exert finer control over the elements in your tensors.\n",
        "\n",
        "* Check out the slicing ops available with TensorFlow NumPy such as `tf.experimental.numpy.take_along_axis` and `tf.experimental.numpy.take`.\n",
        "\n",
        "* Also check out the [Tensor guide](https://www.tensorflow.org/guide/tensor) and the [Variable guide](https://www.tensorflow.org/guide/variable)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "tensor_slicing.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
